"""storage/index_page.py
B+Tree 索引页的最低层包装类。
依赖公共 Page 封装，直接在 Page.data 上读写，不复制内存。
目前仅支持 INT(4B) 键，变长键后续扩展。
页格式 (offset -> meaning):
0   : 1B  page_type (0 = internal, 1 = leaf)
1   : 1B  reserved
2   : 2B  key_count (unsigned short)
4   : 4B  right_sibling_page_id (仅 leaf 页有效; internal 页保留)
8   : 开始 Key/Ptr 区

内部页条目:  (key:int, child_page_id:int)  共8B
叶子页条目: (key:int, record_page_id:int, slot:int) 共12B

最大容量根据 PAGE_SIZE(4096) 自动推算。
"""

import struct
from .page import Page, PAGE_SIZE
import bisect

# constants
PAGE_TYPE_INTERNAL = 0
PAGE_TYPE_LEAF = 1
HEADER_FORMAT = "BBHII"  # page_type, reserved, key_count, right_sibling, parent_page
HEADER_SIZE = struct.calcsize(HEADER_FORMAT)  # 1+1+2+4+4 =12

# entry formats
INTERNAL_ENTRY_FORMAT = "!II"  # key, child_page_id
INTERNAL_ENTRY_SIZE = struct.calcsize(INTERNAL_ENTRY_FORMAT)

LEAF_ENTRY_FORMAT = "!III"  # key, record_page_id, slot_id
LEAF_ENTRY_SIZE = struct.calcsize(LEAF_ENTRY_FORMAT)

class IndexPage:
    def __init__(self, page: Page):
        self.page = page
        self.data = page.data
        # If fresh page (zero), default init as leaf with zero keys
        page_type, _, key_cnt, _ = struct.unpack_from(HEADER_FORMAT, self.data, 0)
        if page_type not in (PAGE_TYPE_INTERNAL, PAGE_TYPE_LEAF):
            self._init_page(PAGE_TYPE_LEAF)  # default leaf

    # ------------- internal helpers -------------
    def _init_page(self, page_type:int, parent:int=0, right_sib:int=0):
        struct.pack_into(HEADER_FORMAT, self.data, 0, page_type, 0, 0, right_sib, parent)
        self.page.is_dirty = True

    def _read_header(self):
        return struct.unpack_from(HEADER_FORMAT, self.data, 0)

    def _write_header(self, page_type:int, key_count:int, right_sibling:int, parent:int=None):
        _, _, _, _, cur_parent = self._read_header()
        if parent is None:
            parent = cur_parent
        struct.pack_into(HEADER_FORMAT, self.data, 0, page_type, 0, key_count, right_sibling, parent)
        self.page.is_dirty = True

    def get_parent(self):
        return struct.unpack_from("!I", self.data, 8)[0]

    def set_parent(self, pid:int):
        page_type,res,cnt,sib,_=struct.unpack_from(HEADER_FORMAT,self.data,0)
        struct.pack_into(HEADER_FORMAT,self.data,0,page_type,res,cnt,sib,pid)
        self.page.is_dirty=True

    # ------------- public APIs -------------
    @property
    def is_leaf(self):
        page_type, _, __, ___ = self._read_header()
        return page_type == PAGE_TYPE_LEAF

    @property
    def key_count(self):
        _, _, cnt, _ = self._read_header()
        return cnt

    def right_sibling(self):
        _, _, _, sib = self._read_header()
        return sib

    def iter_entries(self):
        """Yield tuples representing each entry."""
        cnt = self.key_count
        offset = HEADER_SIZE
        for _ in range(cnt):
            if self.is_leaf:
                key, rpage, slot = struct.unpack_from(LEAF_ENTRY_FORMAT, self.data, offset)
                yield key, (rpage, slot)
                offset += LEAF_ENTRY_SIZE
            else:
                key, child = struct.unpack_from(INTERNAL_ENTRY_FORMAT, self.data, offset)
                yield key, child
                offset += INTERNAL_ENTRY_SIZE

    def _get_keys(self):
        """Return list of keys in this page (for binary search)."""
        keys = []
        offset = HEADER_SIZE
        for _ in range(self.key_count):
            key = struct.unpack_from("!I", self.data, offset)[0]
            keys.append(key)
            offset += LEAF_ENTRY_SIZE if self.is_leaf else INTERNAL_ENTRY_SIZE
        return keys

    def _binary_search(self, key:int):
        keys = self._get_keys()
        return bisect.bisect_left(keys, key)

    def _shift_and_insert(self, entry_bytes:bytes, pos:int, entry_size:int):
        cnt = self.key_count
        if cnt * entry_size + HEADER_SIZE + entry_size > PAGE_SIZE:
            return False
        insert_offset = HEADER_SIZE + pos * entry_size
        tail_offset = HEADER_SIZE + cnt * entry_size
        # shift tail right
        self.data[insert_offset+entry_size:tail_offset+entry_size] = self.data[insert_offset:tail_offset]
        # write new entry
        self.data[insert_offset:insert_offset+entry_size] = entry_bytes
        self._write_header(PAGE_TYPE_LEAF if self.is_leaf else PAGE_TYPE_INTERNAL, cnt+1, self.right_sibling())
        self.page.is_dirty=True
        return True

    def insert_leaf_entry(self, key:int, rid:tuple):
        if not self.is_leaf:
            raise ValueError("Not leaf")
        rpage, slot = rid
        pos = self._binary_search(key)
        entry_bytes = struct.pack(LEAF_ENTRY_FORMAT, key, rpage, slot)
        return self._shift_and_insert(entry_bytes, pos, LEAF_ENTRY_SIZE)

    def insert_internal_entry(self, key:int, child_page_id:int):
        if self.is_leaf:
            raise ValueError("Not internal")
        pos = self._binary_search(key)
        entry_bytes = struct.pack(INTERNAL_ENTRY_FORMAT, key, child_page_id)
        return self._shift_and_insert(entry_bytes, pos, INTERNAL_ENTRY_SIZE)

    def set_right_sibling(self, page_id:int):
        page_type, res, cnt, _ = self._read_header()
        struct.pack_into(HEADER_FORMAT, self.data, 0, page_type, res, cnt, page_id)
        self.page.is_dirty = True

    # splitting helper (simplistic)
    def split_leaf(self, buffer_manager):
        """Split current leaf into two, return (new_page, split_key)."""
        if not self.is_leaf:
            raise ValueError("split only leaf")
        cnt = self.key_count
        mid = cnt // 2
        # Create new page
        new_page_obj = buffer_manager.new_page()
        new_index_page = IndexPage(new_page_obj)
        new_index_page._init_page(PAGE_TYPE_LEAF)
        new_index_page.set_parent(self.get_parent())
        # move second half entries
        src_offset = HEADER_SIZE + mid * LEAF_ENTRY_SIZE
        dst_offset = HEADER_SIZE
        bytes_to_move = (cnt - mid) * LEAF_ENTRY_SIZE
        new_index_page.data[dst_offset:dst_offset+bytes_to_move] = self.data[src_offset:src_offset+bytes_to_move]
        new_index_page._write_header(PAGE_TYPE_LEAF, cnt - mid, self.right_sibling())
        # trim original
        self._write_header(PAGE_TYPE_LEAF, mid, new_page_obj.page_id)
        self.page.is_dirty = True
        new_index_page.page.is_dirty = True
        # split key = first key of new page
        first_key, _ = next(new_index_page.iter_entries())
        return new_index_page, first_key

    def split_internal(self, buffer_manager):
        """Split current internal page, return (new_page_index, split_key)."""
        if self.is_leaf:
            raise ValueError("split_internal called on leaf")
        cnt = self.key_count
        mid = cnt // 2
        new_page_obj = buffer_manager.new_page()
        new_idx_page = IndexPage(new_page_obj)
        new_idx_page._init_page(PAGE_TYPE_INTERNAL)
        new_idx_page.set_parent(self.get_parent())
        # Move entries mid+1 .. cnt-1 (keep mid as split_key) to new page
        src_offset = HEADER_SIZE + (mid + 1) * INTERNAL_ENTRY_SIZE
        entries_to_move = cnt - mid - 1
        if entries_to_move > 0:
            dst_offset = HEADER_SIZE
            bytes_to_move = entries_to_move * INTERNAL_ENTRY_SIZE
            new_idx_page.data[dst_offset:dst_offset + bytes_to_move] = self.data[src_offset:src_offset + bytes_to_move]
        # Update headers
        # split_key entry value
        split_key, child_after_split = struct.unpack_from(INTERNAL_ENTRY_FORMAT, self.data, HEADER_SIZE + mid * INTERNAL_ENTRY_SIZE)
        new_idx_page._write_header(PAGE_TYPE_INTERNAL, entries_to_move, 0)
        self._write_header(PAGE_TYPE_INTERNAL, mid + 1, 0)  # keep split_key in left page (can also remove but easier)
        self.page.is_dirty = True
        new_idx_page.page.is_dirty = True
        return new_idx_page, split_key
